import pymysql
from pymysql.cursors import DictCursor
import getpass
import sys
import os

# 将项目根目录添加到Python路径
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

def get_mysql_credentials():
    """获取用户输入的MySQL连接信息"""
    print("\n====== MySQL数据库配置 ======")
    print("请根据您的MySQL服务器配置输入以下信息：")
    
    host = input("MySQL主机地址 [localhost]: ") or "localhost"
    port = input("MySQL端口 [3306]: ") or "3306"
    user = input("MySQL用户名 [root]: ") or "root"
    password = getpass.getpass("MySQL密码 (输入时不会显示): ")
    db_name = input("数据库名称 [mobile_test_platform]: ") or "mobile_test_platform"
    
    return {
        'host': host,
        'port': int(port),
        'user': user,
        'password': password,
        'db_name': db_name
    }

def create_database_and_tables(config):
    """创建数据库和表结构"""
    try:
        # 先尝试不指定数据库连接
        conn = pymysql.connect(
            host=config['host'],
            port=config['port'],
            user=config['user'],
            password=config['password'],
            charset='utf8mb4',
            cursorclass=DictCursor
        )
        cursor = conn.cursor()
        
        print(f"\n正在连接MySQL服务器: {config['host']}:{config['port']}...")
        
        # 创建数据库
        print(f"正在创建数据库 {config['db_name']}...")
        create_db_sql = f"CREATE DATABASE IF NOT EXISTS {config['db_name']} DEFAULT CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;"
        cursor.execute(create_db_sql)
        conn.commit()
        
        # 选择创建的数据库
        cursor.execute(f"USE {config['db_name']};")
        
        # 创建表结构
        print("正在创建表结构...")
        
        # 创建devices表
        create_devices_table = '''
        CREATE TABLE IF NOT EXISTS devices (
            id INT PRIMARY KEY AUTO_INCREMENT,
            device_id VARCHAR(100) NOT NULL UNIQUE,
            device_name VARCHAR(200),
            model VARCHAR(200),
            brand VARCHAR(100),
            os_version VARCHAR(50),
            status VARCHAR(50) DEFAULT 'offline',
            ip_address VARCHAR(50),
            port INT,
            battery_level INT DEFAULT 0,
            is_usb_connected BOOLEAN DEFAULT TRUE,
            current_task_id INT,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
            FOREIGN KEY (current_task_id) REFERENCES tasks(id)
        );
        '''
        cursor.execute(create_devices_table)
        
        # 创建scripts表
        create_scripts_table = '''
        CREATE TABLE IF NOT EXISTS scripts (
            id INT PRIMARY KEY AUTO_INCREMENT,
            name VARCHAR(200) NOT NULL,
            description TEXT,
            content TEXT NOT NULL,
            file_path VARCHAR(255),
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP
        );
        '''
        cursor.execute(create_scripts_table)
        
        # 创建tasks表
        create_tasks_table = '''
        CREATE TABLE IF NOT EXISTS tasks (
            id INT PRIMARY KEY AUTO_INCREMENT,
            name VARCHAR(200) NOT NULL,
            description TEXT,
            script_id INT NOT NULL,
            status VARCHAR(50) DEFAULT 'pending',
            priority INT DEFAULT 1,
            start_time DATETIME,
            end_time DATETIME,
            created_by VARCHAR(100),
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            updated_at DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
            FOREIGN KEY (script_id) REFERENCES scripts(id)
        );
        '''
        cursor.execute(create_tasks_table)
        
        conn.commit()
        
        print("\n✅ 数据库初始化成功！")
        print(f"数据库: {config['db_name']}")
        print(f"主机: {config['host']}:{config['port']}")
        print(f"用户: {config['user']}")
        
        # 更新app.py中的数据库连接配置
        update_app_config(config)
        
        return True
        
    except pymysql.MySQLError as e:
        print(f"\n❌ MySQL操作失败: {e}")
        print("请检查以下几点：")
        print("1. MySQL服务器是否正在运行")
        print("2. 用户名和密码是否正确")
        print("3. 用户是否有足够的权限创建数据库和表")
        print("4. 如果使用root用户，请确保配置了正确的认证方式")
        return False
    finally:
        # 关闭连接
        if 'cursor' in locals():
            cursor.close()
        if 'conn' in locals():
            conn.close()

def update_app_config(config):
    """更新app.py中的数据库连接配置"""
    try:
        with open('app.py', 'r', encoding='utf-8') as f:
            content = f.read()
        
        # 构建新的数据库连接字符串
        new_db_uri = f"mysql+pymysql://{config['user']}:{config['password']}@{config['host']}:{config['port']}/{config['db_name']}"
        
        # 替换数据库连接配置
        import re
        
        # 使用正则表达式替换数据库连接配置
        content = re.sub(
            r'app.config\[\'SQLALCHEMY_DATABASE_URI\'\]\s*=\s*[\'"][^\'"]*[\'"]',
            f"app.config['SQLALCHEMY_DATABASE_URI'] = '{new_db_uri}'",
            content
        )
        
        with open('app.py', 'w', encoding='utf-8') as f:
            f.write(content)
        
        print("\n✅ app.py中的数据库配置已更新！")
        
    except Exception as e:
        print(f"\n⚠️ 更新app.py配置时出错: {e}")
        print("请手动更新app.py中的数据库连接配置。")
        print(f"建议的配置: app.config['SQLALCHEMY_DATABASE_URI'] = '{new_db_uri}'")

def main():
    print("====== 移动测试平台数据库初始化工具 ======")
    print("此工具将帮助您配置MySQL数据库连接并初始化数据库结构。")
    
    while True:
        config = get_mysql_credentials()
        if create_database_and_tables(config):
            break
        
        retry = input("\n是否重试？(y/n): ")
        if retry.lower() != 'y':
            print("\n数据库初始化已取消。")
            break
    
    print("\n====== 操作完成 ======")

if __name__ == "__main__":
    main()